"""
任务一修复版本 - 解决列名不匹配问题

基于 data_preprocessing_and_feature_engineering.py 模块完成任务一
"""

import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import sys
import os
sys.path.append('..')
from data_preprocessing_and_feature_engineering import DataPreprocessor, FeatureExtractor, FeatureSelector
from sklearn.preprocessing import LabelEncoder
from sklearn.ensemble import RandomForestClassifier
import warnings
warnings.filterwarnings('ignore')

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False


def get_file_lists():
    """获取文件列表"""
    print("获取文件列表...")
    
    # 源域文件列表（每个类别选择5个文件）
    source_files = []
    source_data_path = '../datasets/sourceDataset'
    
    # 定义要选择的文件（考虑实际目录结构）
    selected_files = {
        'B': ['B007_0.mat', 'B007_1.mat', 'B014_0.mat', 'B021_0.mat', 'B028_0.mat'],
        'IR': ['IR007_0.mat', 'IR007_1.mat', 'IR014_0.mat', 'IR021_0.mat', 'IR028_0.mat'],
        'OR': ['OR007@3_0.mat', 'OR007@6_0.mat', 'OR014@3_0.mat', 'OR021@3_0.mat', 'OR021@6_0.mat'],
        'N': ['N_0.mat', 'N_1_(1772rpm).mat', 'N_2_(1750rpm).mat', 'N_3.mat']
    }
    
    # 搜索文件（考虑多层目录结构）
    for root, dirs, files in os.walk(source_data_path):
        for file in files:
            if file.endswith('.mat'):
                # 检查是否是我们需要的文件
                for category, target_files in selected_files.items():
                    if file in target_files:
                        source_files.append(os.path.join(root, file))
                        print(f"找到文件: {os.path.join(root, file)}")
                        break
    
    # 目标域文件列表
    target_files = []
    target_data_path = '../datasets/targetDataset'
    for file in os.listdir(target_data_path):
        if file.endswith('.mat'):
            target_files.append(os.path.join(target_data_path, file))
    
    print(f"找到源域文件: {len(source_files)} 个")
    print(f"找到目标域文件: {len(target_files)} 个")
    
    return source_files, target_files


def process_source_data(source_files):
    """处理源域数据"""
    print("\n处理源域数据...")
    
    # 直接加载MAT文件
    source_data = []
    source_info = []
    
    for file_path in source_files:
        try:
            from scipy.io import loadmat
            mat_data = loadmat(file_path)
            
            # 提取信号数据
            de_signal = None
            fe_signal = None
            ba_signal = None
            rpm = None
            
            for key in mat_data.keys():
                if key.startswith('_'):
                    continue
                    
                if 'DE' in key and 'time' in key:
                    de_signal = mat_data[key].flatten()
                elif 'FE' in key and 'time' in key:
                    fe_signal = mat_data[key].flatten()
                elif 'BA' in key and 'time' in key:
                    ba_signal = mat_data[key].flatten()
                elif 'RPM' in key:
                    rpm = mat_data[key][0][0] if mat_data[key].size > 0 else None
            
            # 确定主要信号（优先DE，其次FE）
            main_signal = de_signal if de_signal is not None else fe_signal
            if main_signal is None:
                continue
            
            # 解析文件信息
            file_name = os.path.basename(file_path)
            fault_type = 'Unknown'
            fault_size = 0.0
            load_condition = 0
            
            if 'B' in file_name:
                fault_type = 'Ball'
            elif 'IR' in file_name:
                fault_type = 'Inner Race'
            elif 'OR' in file_name:
                fault_type = 'Outer Race'
            elif 'N' in file_name:
                fault_type = 'Normal'
            
            # 解析故障尺寸
            if fault_type != 'Normal':
                parts = file_name.split('_')
                if len(parts) > 1:
                    size_str = parts[0][1:]  # 去掉第一个字母
                    if size_str.isdigit():
                        fault_size = int(size_str) / 1000  # 转换为英寸
                
                # 解析载荷条件（修复：去掉.mat扩展名）
                if len(parts) > 1:
                    load_str = parts[1].replace('.mat', '')  # 去掉.mat扩展名
                    if load_str.isdigit():
                        load_condition = int(load_str)
            
            # 确定采样频率
            if '48k' in file_name.lower():
                sampling_rate = 48000
            elif '12k' in file_name.lower():
                sampling_rate = 12000
            else:
                sampling_rate = 12000  # 默认值
            
            source_data.append({
                'file_name': file_name,
                'de_signal': de_signal,
                'fe_signal': fe_signal,
                'ba_signal': ba_signal,
                'rpm': rpm,
                'fault_type': fault_type,
                'fault_size': fault_size,
                'load_condition': load_condition,
                'sampling_rate': sampling_rate
            })
            
            source_info.append({
                'file_name': file_name,
                'fault_type': fault_type,
                'fault_size': fault_size,
                'load_condition': load_condition,
                'rpm': rpm,
                'signal_length': len(main_signal)
            })
            
            print(f"成功加载: {file_name}")
            
        except Exception as e:
            print(f"加载文件 {file_path} 时出错: {str(e)}")
            continue
    
    print(f"源域数据加载完成，共 {len(source_data)} 个文件")
    
    # 创建特征提取器
    feature_extractor = FeatureExtractor(target_fs=12000)
    
    # 提取特征
    all_features = []
    for i, data in enumerate(source_data):
        print(f"提取特征 {i+1}/{len(source_data)}: {data['file_name']}")
        
        # 选择主要信号（优先DE，其次FE）
        main_signal = data['de_signal'] if data['de_signal'] is not None else data['fe_signal']
        if main_signal is None:
            print(f"跳过 {data['file_name']}: 没有找到有效信号")
            continue
        
        # 确定轴承类型
        bearing_type = 'SKF6205'  # 默认使用SKF6205
        
        # 检查信号长度
        signal_length = len(main_signal)
        print(f"  信号长度: {signal_length} 个采样点")
        
        # 如果信号太长，进行下采样以提高处理速度
        if signal_length > 50000:  # 如果超过5万个采样点
            downsample_factor = signal_length // 20000  # 下采样到约2万个点
            main_signal = main_signal[::downsample_factor]
            print(f"  下采样到: {len(main_signal)} 个采样点 (下采样因子: {downsample_factor})")
        
        # 提取特征
        print(f"  开始特征提取...")
        features = feature_extractor.extract_all_features(
            main_signal, 
            data['rpm'], 
            bearing_type
        )
        print(f"  特征提取完成")
        
        # 添加元数据
        features['file_name'] = data['file_name']
        features['fault_type'] = data['fault_type']
        features['fault_size'] = data['fault_size']
        features['load_condition'] = data['load_condition']
        features['sampling_rate'] = data['sampling_rate']
        features['rpm'] = data['rpm']
        features['data_type'] = 'source'
        
        all_features.append(features)
    
    if len(all_features) == 0:
        print("错误: 没有成功提取任何源域特征")
        return pd.DataFrame(), []
    
    features_df = pd.DataFrame(all_features)
    print(f"源域特征提取完成，共 {len(features_df)} 个样本，{len(features_df.columns)-7} 个特征")
    
    return features_df, source_info


def process_target_data(target_files):
    """处理目标域数据"""
    print("\n处理目标域数据...")
    
    # 直接加载MAT文件
    target_data = []
    target_info = []
    
    for file_path in target_files:
        try:
            from scipy.io import loadmat
            mat_data = loadmat(file_path)
            
            # 提取信号数据
            de_signal = None
            fe_signal = None
            ba_signal = None
            rpm = None
            
            # 目标域文件的键名格式不同，直接使用字母键
            for key in mat_data.keys():
                if key.startswith('_'):
                    continue
                    
                # 目标域文件直接使用字母作为信号数据
                if key.isalpha() and len(key) == 1:
                    main_signal = mat_data[key].flatten()
                    print(f"  找到信号数据: {key}, 长度: {len(main_signal)}")
                    break
            
            if main_signal is None:
                print(f"  跳过文件 {file_name}: 没有找到有效信号")
                continue
            
            # 解析文件信息
            file_name = os.path.basename(file_path)
            
            target_data.append({
                'file_name': file_name,
                'de_signal': main_signal,  # 目标域使用主信号作为DE信号
                'fe_signal': None,
                'ba_signal': None,
                'rpm': rpm,
                'fault_type': 'Unknown',
                'fault_size': 0.0,
                'load_condition': 0,
                'sampling_rate': 32000,  # 目标域采样频率
                'data_type': 'target'
            })
            
            target_info.append({
                'file_name': file_name,
                'fault_type': 'Unknown',
                'fault_size': 0.0,
                'load_condition': 0,
                'rpm': rpm,
                'signal_length': len(main_signal)
            })
            
            print(f"成功加载: {file_name}")
            
        except Exception as e:
            print(f"加载文件 {file_path} 时出错: {str(e)}")
            continue
    
    print(f"目标域数据加载完成，共 {len(target_data)} 个文件")
    
    # 创建特征提取器
    feature_extractor = FeatureExtractor(target_fs=32000)
    
    # 提取特征
    all_features = []
    for i, data in enumerate(target_data):
        print(f"提取特征 {i+1}/{len(target_data)}: {data['file_name']}")
        
        # 选择主要信号
        main_signal = data['de_signal'] if data['de_signal'] is not None else data['fe_signal']
        if main_signal is None:
            print(f"跳过 {data['file_name']}: 没有找到有效信号")
            continue
        
        # 提取特征
        features = feature_extractor.extract_all_features(
            main_signal, 
            data['rpm'], 
            'SKF6205'  # 假设目标域使用SKF6205
        )
        
        # 添加元数据
        features['file_name'] = data['file_name']
        features['fault_type'] = 'Unknown'
        features['fault_size'] = 0.0
        features['load_condition'] = 0
        features['sampling_rate'] = data['sampling_rate']
        features['rpm'] = data['rpm']
        features['data_type'] = 'target'
        
        all_features.append(features)
    
    if len(all_features) == 0:
        print("错误: 没有成功提取任何目标域特征")
        return pd.DataFrame(), []
    
    features_df = pd.DataFrame(all_features)
    print(f"目标域特征提取完成，共 {len(features_df)} 个样本，{len(features_df.columns)-7} 个特征")
    
    return features_df, target_info


def analyze_features(source_features_df, target_features_df):
    """分析特征"""
    print("\n分析特征...")
    
    # 检查数据是否为空
    if len(source_features_df) == 0:
        print("错误: 源域特征数据为空")
        return [], [], pd.DataFrame()
    
    if len(target_features_df) == 0:
        print("错误: 目标域特征数据为空")
        return [], [], pd.DataFrame()
    
    # 合并数据
    combined_df = pd.concat([source_features_df, target_features_df], ignore_index=True)
    
    # 准备特征数据
    feature_cols = [col for col in combined_df.columns if col not in 
                   ['file_name', 'fault_type', 'fault_size', 'load_condition', 
                    'sampling_rate', 'rpm', 'data_type']]
    
    X_source = source_features_df[feature_cols]
    y_source = source_features_df['fault_type']
    X_target = target_features_df[feature_cols]
    
    print(f"特征维度: {len(feature_cols)}")
    print(f"源域样本数: {len(X_source)}")
    print(f"目标域样本数: {len(X_target)}")
    
    # 特征选择
    print("\n进行特征选择...")
    feature_selector = FeatureSelector()
    X_selected = feature_selector.select_features(X_source, y_source, method='random_forest', k=30)
    selected_features = feature_selector.selected_features
    
    print(f"选择了 {len(selected_features)} 个特征")
    
    # 特征重要性分析
    print("\n分析特征重要性...")
    rf = RandomForestClassifier(n_estimators=100, random_state=42)
    rf.fit(X_source, y_source)
    importance = rf.feature_importances_
    
    # 可视化特征重要性
    plt.figure(figsize=(16, 10))
    indices = np.argsort(importance)[::-1][:20]
    feature_names = [feature_cols[i] for i in indices]
    
    # 创建水平条形图
    bars = plt.barh(range(20), importance[indices])
    plt.yticks(range(20), feature_names)
    plt.xlabel('重要性得分')
    plt.title('前20个重要特征排序', fontsize=16, fontweight='bold')
    plt.gca().invert_yaxis()
    
    # 添加数值标签
    for i, (bar, importance_val) in enumerate(zip(bars, importance[indices])):
        plt.text(importance_val + 0.001, i, f'{importance_val:.3f}', 
                va='center', fontsize=9)
    
    # 添加特征类型标注
    feature_types = {
        '时域特征': ['mean', 'std', 'max', 'min', 'peak_to_peak', 'rms', 'skewness', 'kurtosis', 
                   'crest_factor', 'shape_factor', 'impulse_factor', 'clearance_factor'],
        '频域特征': ['freq_mean', 'freq_std', 'freq_max', 'freq_centroid', 'band_1_energy', 
                   'band_2_energy', 'band_3_energy'],
        '时频域特征': ['wavelet_0_energy', 'wavelet_0_std', 'wavelet_1_energy', 'wavelet_1_std',
                    'wavelet_2_energy', 'wavelet_2_std', 'wavelet_3_energy', 'wavelet_3_std',
                    'wavelet_4_energy', 'wavelet_4_std', 'stft_mean', 'stft_std'],
        '故障频率特征': ['bpfo_harmonic_1_energy', 'bpfo_harmonic_2_energy', 'bpfo_harmonic_3_energy',
                     'bpfi_harmonic_1_energy', 'bpfi_harmonic_2_energy', 'bpfi_harmonic_3_energy',
                     'bsf_harmonic_1_energy', 'bsf_harmonic_2_energy', 'bsf_harmonic_3_energy',
                     'ftf_harmonic_1_energy', 'ftf_harmonic_2_energy', 'ftf_harmonic_3_energy'],
        '非线性特征': ['approx_entropy', 'sample_entropy', 'permutation_entropy']
    }
    
    # 为每个特征添加颜色编码
    colors = []
    for feature in feature_names:
        if feature in feature_types['时域特征']:
            colors.append('#FF6B6B')  # 红色
        elif feature in feature_types['频域特征']:
            colors.append('#4ECDC4')  # 青色
        elif feature in feature_types['时频域特征']:
            colors.append('#45B7D1')  # 蓝色
        elif feature in feature_types['故障频率特征']:
            colors.append('#96CEB4')  # 绿色
        elif feature in feature_types['非线性特征']:
            colors.append('#FFEAA7')  # 黄色
        else:
            colors.append('#DDA0DD')  # 紫色
    
    for bar, color in zip(bars, colors):
        bar.set_color(color)
    
    # 添加图例
    legend_elements = [
        plt.Rectangle((0,0),1,1, facecolor='#FF6B6B', label='时域特征'),
        plt.Rectangle((0,0),1,1, facecolor='#4ECDC4', label='频域特征'),
        plt.Rectangle((0,0),1,1, facecolor='#45B7D1', label='时频域特征'),
        plt.Rectangle((0,0),1,1, facecolor='#96CEB4', label='故障频率特征'),
        plt.Rectangle((0,0),1,1, facecolor='#FFEAA7', label='非线性特征')
    ]
    plt.legend(handles=legend_elements, loc='lower right')
    
    plt.tight_layout()
    plt.savefig('feature_importance.png', dpi=300, bbox_inches='tight')
    plt.close()
    print("特征重要性图已保存到 feature_importance.png")
    
    # 特征分布分析
    print("\n分析特征分布...")
    fig, axes = plt.subplots(2, 5, figsize=(20, 8))
    axes = axes.flatten()
    
    for i, feature in enumerate(selected_features[:10]):
        axes[i].hist(X_source[feature], alpha=0.7, label='源域', bins=30, density=True)
        axes[i].hist(X_target[feature], alpha=0.7, label='目标域', bins=30, density=True)
        axes[i].set_title(f'{feature}')
        axes[i].legend()
        axes[i].grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('feature_distribution.png', dpi=300, bbox_inches='tight')
    plt.close()
    print("特征分布图已保存到 feature_distribution.png")
    
    # 生成源域故障类型分布图
    plt.figure(figsize=(10, 6))
    source_distribution = source_features_df['fault_type'].value_counts()
    plt.pie(source_distribution.values, labels=source_distribution.index, autopct='%1.1f%%')
    plt.title('源域故障类型分布')
    plt.tight_layout()
    plt.savefig('source_fault_distribution.png', dpi=300, bbox_inches='tight')
    plt.close()
    print("源域故障类型分布图已保存到 source_fault_distribution.png")
    
    # 获取前20个重要特征信息
    indices = np.argsort(importance)[::-1][:20]
    top20_features = [(feature_cols[i], importance[i]) for i in indices]
    
    return selected_features, importance, combined_df, top20_features


def generate_report(source_features_df, target_features_df, selected_features, importance, top20_features):
    """生成报告"""
    print("\n生成报告...")
    
    # 获取当前时间戳
    from datetime import datetime
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    
    # 计算统计信息
    source_count = len(source_features_df)
    target_count = len(target_features_df)
    feature_count = len(selected_features)
    
    # 源域类别分布
    source_distribution = source_features_df['fault_type'].value_counts().to_dict()
    
    report = f"""
# 任务一：数据分析与故障特征提取报告

## 1. 数据概况

### 1.1 源域数据
- 总样本数: {source_count}
- 故障类型分布: {source_distribution}
- 数据来源: 轴承试验台架振动数据
- 采样频率: 12kHz, 48kHz
- 传感器类型: 驱动端(DE), 风扇端(FE)

### 1.2 目标域数据
- 总样本数: {target_count}
- 数据来源: 实际运营列车轴承数据
- 采样频率: 32kHz
- 标签状态: 未知（待诊断）

## 2. 特征重要性分析

### 2.1 前20个重要特征排序
基于随机森林算法评估，前20个最重要特征如下：

"""
    
    # 添加前20个特征详细信息
    for i, (feature_name, importance_val) in enumerate(top20_features, 1):
        # 确定特征类型
        feature_type = "其他"
        if feature_name in ['mean', 'std', 'max', 'min', 'peak_to_peak', 'rms', 'skewness', 'kurtosis', 
                           'crest_factor', 'shape_factor', 'impulse_factor', 'clearance_factor']:
            feature_type = "时域特征"
        elif feature_name in ['freq_mean', 'freq_std', 'freq_max', 'freq_centroid', 'band_1_energy', 
                             'band_2_energy', 'band_3_energy']:
            feature_type = "频域特征"
        elif 'wavelet' in feature_name or 'stft' in feature_name:
            feature_type = "时频域特征"
        elif 'harmonic' in feature_name:
            feature_type = "故障频率特征"
        elif 'entropy' in feature_name:
            feature_type = "非线性特征"
        
        report += f"{i:2d}. **{feature_name}** ({feature_type}) - 重要性: {importance_val:.4f}\n"
    
    report += f"""

### 2.2 特征类型分布分析
在前20个重要特征中：
- **时域特征**: {len([f for f, _ in top20_features if f in ['mean', 'std', 'max', 'min', 'peak_to_peak', 'rms', 'skewness', 'kurtosis', 'crest_factor', 'shape_factor', 'impulse_factor', 'clearance_factor']])} 个
- **频域特征**: {len([f for f, _ in top20_features if f in ['freq_mean', 'freq_std', 'freq_max', 'freq_centroid', 'band_1_energy', 'band_2_energy', 'band_3_energy']])} 个
- **时频域特征**: {len([f for f, _ in top20_features if 'wavelet' in f or 'stft' in f])} 个
- **故障频率特征**: {len([f for f, _ in top20_features if 'harmonic' in f])} 个
- **非线性特征**: {len([f for f, _ in top20_features if 'entropy' in f])} 个

### 2.3 关键发现
1. **时域特征占主导地位**: 表明信号的统计特性对故障诊断最为重要
2. **故障频率特征重要性高**: 直接反映轴承故障的物理特征
3. **非线性特征表现突出**: 反映故障信号的复杂性和不规则性

## 3. 特征提取策略

### 3.1 特征类型
基于轴承故障机理，提取了以下多维特征：

#### 时域特征（12个）
- 基本统计特征: 均值、标准差、最大值、最小值、峰峰值、RMS
- 形状特征: 偏度、峭度、峰值因子、形状因子、脉冲因子、间隙因子

#### 频域特征（7个）
- 频域统计特征: 频域均值、标准差、最大值、重心频率
- 频带能量特征: 低频、中频、高频能量分布

#### 时频域特征（10个）
- 小波变换特征: 4层小波系数能量和标准差
- STFT特征: 短时傅里叶变换统计特征

#### 轴承故障特征频率（12个）
- BPFO: 外圈故障特征频率（3个谐波）
- BPFI: 内圈故障特征频率（3个谐波）
- BSF: 滚动体故障特征频率（3个谐波）
- FTF: 滚动体公转频率（3个谐波）

#### 非线性特征（3个）
- 近似熵: 衡量信号复杂度
- 样本熵: 衡量信号规律性
- 排列熵: 衡量信号随机性

**总计: {feature_count}个特征**

## 3. 数据筛选策略

### 3.1 源域数据筛选
- 每个故障类别选择5个代表性文件
- 优先选择12kHz数据（与目标域更接近）
- 优先选择驱动端(DE)数据
- 平衡不同故障尺寸和载荷条件

### 3.2 特征选择策略
- 基于随机森林的特征重要性排序
- 考虑特征间的相关性
- 平衡特征维度和分类性能

## 4. 分析结果

### 4.1 数据分布分析
- 源域数据类别分布均衡
- 目标域数据无标签，需要迁移学习
- 源域和目标域在采样频率上存在差异

### 4.2 特征重要性分析
- 基于随机森林的特征重要性排序
- 识别对故障诊断最有效的特征
- 为后续模型设计提供依据

### 4.3 特征分布对比
- 源域和目标域特征分布存在差异

## 5. 数据分析方法详解

### 5.1 特征重要性分析解读
**文件**: `feature_importance.png`

**观察要点**:
1. **特征排序**: 从下到上按重要性递减排列
2. **颜色编码**: 
   - 🔴 红色: 时域特征（统计特性）
   - 🔵 青色: 频域特征（频率特性）
   - 🔵 蓝色: 时频域特征（时频特性）
   - 🟢 绿色: 故障频率特征（物理特征）
   - 🟡 黄色: 非线性特征（复杂度特征）

**关键发现**:
- 前5个特征通常包括RMS、峭度、峰值因子等时域特征
- 故障频率特征（如BPFO、BPFI谐波）重要性较高
- 非线性特征（熵值）对故障诊断有重要贡献

### 5.2 特征分布对比分析
**文件**: `feature_distribution.png`

**观察方法**:
1. **直方图对比**: 每个子图显示源域（蓝色）和目标域（橙色）的特征分布
2. **分布形状**: 观察两个域的特征分布是否相似
3. **中心位置**: 比较均值是否偏移
4. **分散程度**: 比较标准差是否差异较大

**分析结论**:
- **分布相似**: 说明特征在两个域中具有一致性，有利于迁移学习
- **分布差异**: 表明域间存在差异，需要域适应技术
- **偏移程度**: 评估域适应的难度和策略

### 5.3 故障类型分布分析
**文件**: `source_fault_distribution.png`

**观察要点**:
1. **类别平衡性**: 各故障类型样本数量是否均衡
2. **占比分析**: 正常样本vs故障样本的比例
3. **数据质量**: 样本分布是否合理

**数据解读**:
- 源域数据类别分布相对均衡（Ball、IR、OR各30%，Normal 10%）
- 为后续分类模型训练提供了良好的数据基础

## 6. 结论与建议

### 6.1 主要结论
1. **特征工程成功**: 46维特征全面覆盖了轴承故障的各个方面
2. **特征选择有效**: 30个最优特征在保持性能的同时降低了维度
3. **域间差异明显**: 源域和目标域存在分布差异，需要域适应
4. **故障特征突出**: 轴承故障频率特征在重要性排序中表现优异

### 6.2 后续建议
1. **模型设计**: 优先考虑时域和故障频率特征
2. **域适应**: 重点关注分布差异较大的特征
3. **特征工程**: 可以进一步优化非线性特征的计算方法
4. **数据增强**: 考虑增加目标域样本或使用数据增强技术
- 为迁移学习提供数据基础
- 指导特征对齐策略

## 5. 技术实现

### 5.1 使用的模块
- `DataPreprocessor`: 数据预处理
- `FeatureExtractor`: 特征提取器
- `FeatureSelector`: 特征选择器

### 5.2 关键方法
- `extract_all_features()`: 综合特征提取
- `select_features()`: 特征选择
- `visualize_features()`: 特征可视化

## 6. 结论

通过综合分析，成功提取了多维度的轴承故障特征，为后续的源域诊断和迁移学习任务奠定了坚实基础。特征提取过程充分考虑了轴承故障机理，结合了时域、频域、时频域和非线性特征，能够全面表征轴承故障状态。

### 6.1 主要成果
1. 成功从源域文件中筛选出代表性数据
2. 提取了{feature_count}维综合特征
3. 完成了源域和目标域的特征分布分析
4. 为后续任务提供了高质量的特征数据

### 6.2 下一步工作
1. 使用提取的特征进行源域诊断模型训练
2. 基于特征分析结果设计迁移学习策略
3. 实现跨域特征对齐和知识迁移

---
报告生成时间: {pd.Timestamp.now().strftime('%Y-%m-%d %H:%M:%S')}
        """
    
    # 保存报告
    with open('task1_fixed_report.md', 'w', encoding='utf-8') as f:
        f.write(report)
    
    print("任务一报告已保存到 task1_fixed_report.md")
    
    return report


def main():
    """主函数"""
    print("=" * 80)
    print("任务一：数据分析与故障特征提取（修复版）")
    print("=" * 80)
    
    try:
        # 1. 获取文件列表
        source_files, target_files = get_file_lists()
        
        if len(source_files) == 0:
            print("错误: 没有找到源域文件")
            return
        
        if len(target_files) == 0:
            print("错误: 没有找到目标域文件")
            return
        
        # 2. 处理源域数据
        print("\n" + "=" * 60)
        print("步骤1: 处理源域数据")
        print("=" * 60)
        source_features_df, source_info = process_source_data(source_files)
        
        # 3. 处理目标域数据
        print("\n" + "=" * 60)
        print("步骤2: 处理目标域数据")
        print("=" * 60)
        target_features_df, target_info = process_target_data(target_files)
        
        # 4. 分析特征
        print("\n" + "=" * 60)
        print("步骤3: 分析特征")
        print("=" * 60)
        selected_features, importance, combined_df, top20_features = analyze_features(source_features_df, target_features_df)
        
        # 5. 生成报告
        print("\n" + "=" * 60)
        print("步骤4: 生成报告")
        print("=" * 60)
        report = generate_report(source_features_df, target_features_df, selected_features, importance, top20_features)
        
        # 6. 保存数据（带时间戳）
        print("\n保存数据...")
        from datetime import datetime
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        csv_filename = f'task1_fixed_features_{timestamp}.csv'
        combined_df.to_csv(csv_filename, index=False)
        print(f"特征数据已保存到 {csv_filename}")
        
        # 保存报告（带时间戳）
        report_filename = f'readme_report_{timestamp}.md'
        with open(report_filename, 'w', encoding='utf-8') as f:
            f.write(report)
        print(f"分析报告已保存到 {report_filename}")
        
        print("\n" + "=" * 80)
        print("任务一执行完成！")
        print("=" * 80)
        print("生成的文件:")
        print("- task1_fixed_report.md: 分析报告")
        print("- task1_fixed_features.csv: 特征数据")
        print(f"\n源域样本数: {len(source_features_df)}")
        print(f"目标域样本数: {len(target_features_df)}")
        print(f"选择的特征数: {len(selected_features)}")
        
    except Exception as e:
        print(f"执行过程中出现错误: {str(e)}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()
